from typing import Iterable

import torch
from zkl_llmpt_memory import BatchSequenceMemory, VariableMemory


class HaodarTransformerLayerMemory(torch.nn.Module):
    def __init__(self, capacity: int):
        super().__init__()
        self.k_memory = BatchSequenceMemory(capacity=capacity, dim=-3)
        self.v_memory = BatchSequenceMemory(capacity=capacity, dim=-3)

    def get(self) -> tuple[tuple[torch.Tensor, torch.Tensor], torch.Tensor] | tuple[None, None]:
        tokens_k, mask = self.k_memory.get()
        tokens_v, _ = self.v_memory.get()

        if tokens_k is None:
            return None, None

        tokens_kv_mask = mask.unsqueeze(-1).unsqueeze(-1)
        tokens_k = torch.masked_fill(tokens_k, torch.logical_not(tokens_kv_mask), 0.0)
        tokens_v = torch.masked_fill(tokens_v, torch.logical_not(tokens_kv_mask), 0.0)
        return (tokens_k, tokens_v), mask

    def append(self, tokens_kv: tuple[torch.Tensor, torch.Tensor]):
        tokens_k, tokens_v = tokens_kv
        self.k_memory.append(tokens_k)
        self.v_memory.append(tokens_v)

    def reset(self, mask: torch.Tensor | None = None):
        self.k_memory.reset(mask)
        self.v_memory.reset(mask)


class HaodarTransformerMemory(torch.nn.Module):
    def __init__(self, layers_capacity: Iterable[int] | int):
        super().__init__()
        layers_capacity = tuple(layers_capacity) if isinstance(layers_capacity, Iterable) else (layers_capacity,)
        self.layers = tuple(
            HaodarTransformerLayerMemory(layer_capacity) if layer_capacity > 0 else None
            for layer_capacity in layers_capacity)

    def get(self):
        return tuple(zip(*(layer.get() for layer in self.layers if layer is not None)))

    def append(self, layers_tokens_kv: Iterable[tuple[torch.Tensor, torch.Tensor]]):
        for layer, tokens_kv in zip(self.layers, layers_tokens_kv):
            if layer is not None:
                layer.append(tokens_kv)

    def reset(self, mask: torch.Tensor | None = None):
        for layer in self.layers:
            if layer is not None:
                layer.reset(mask)


class RandomOffsetMemory(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.memory = VariableMemory()

    def get(self) -> torch.Tensor | None:
        return self.memory.get()

    def reset(self,
        mask: torch.Tensor,
        period: torch.Tensor | float,
        dtype: torch.dtype | None = None,
        device: torch.device | None = None,
    ):
        value = self.memory.get()
        offset = torch.rand(mask.shape, dtype=dtype, device=device) * period
        value = torch.where(mask, offset, value) \
            if value is not None else offset
        self.memory.set(value)
